import requests
from bs4 import BeautifulSoup
from collections import defaultdict
import re
import json
import argparse


def fetch_pytorch_install_commands(url, command_type='conda'):
    response = requests.get(url)
    response.raise_for_status()

    soup = BeautifulSoup(response.text, 'html.parser')
    pre_blocks = soup.find_all('pre')

    command_map = defaultdict(list)

    for block in pre_blocks:
        lines = block.get_text(strip=True, separator="\n").split('\n')
        for i in range(0, len(lines) - 1, 2):
            if lines[i].startswith('# CUDA') or lines[i].startswith('# CPU Only'):
                command = lines[i + 1].strip()
                if command.startswith(command_type):  # 根据参数过滤命令
                    key = lines[i].strip().lstrip('# ').strip()
                    command_map[key].append(command)

    return command_map


def extract_version(text):
    match = re.search(r'(\d+\.\d+)', text)
    return float(match.group(1)) if match else float('-inf')


def sort_keys(command_map):
    cuda_keys = [key for key in command_map if key.startswith('CUDA')]
    other_keys = [key for key in command_map if not key.startswith('CUDA')]

    cuda_keys.sort(key=extract_version, reverse=True)

    return cuda_keys + other_keys


def save_to_json(data, filename):
    with open(filename, 'w') as f:
        json.dump(data, f, indent=4)


def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(
        description='Fetch PyTorch install commands.')
    parser.add_argument('--command-type', type=str, default='conda',
                        choices=['conda', 'pip'], help='Type of commands to fetch (conda or pip).')
    parser.add_argument('--output', type=str, default='pytorch_versions.json',
                        help='Output JSON filename.')
    args = parser.parse_args()

    url = 'https://pytorch.org/get-started/previous-versions/'
    command_map = fetch_pytorch_install_commands(url, args.command_type)

    sorted_keys = sort_keys(command_map)
    sorted_command_map = {key: command_map[key] for key in sorted_keys}

    save_to_json(sorted_command_map, args.output)

    # Optionally print the JSON data
    print(json.dumps(sorted_command_map, indent=4))


if __name__ == '__main__':
    main()
